#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
直接优化解决方案 - 包含遗传算法思想
"""

import pandas as pd
import numpy as np

def main():
    # 加载数据
    df = pd.read_excel('附件1.xlsx')
    df.columns = ['sequence', 'config', 'color', 'drive']
    
    cars = []
    color_map = {}
    drive_map = {}
    
    for _, row in df.iterrows():
        car_id = int(row['sequence'])
        cars.append(car_id)
        color_map[car_id] = str(row['color'])
        drive_map[car_id] = str(row['drive'])
    
    print(f"加载 {len(cars)} 辆车")
    
    def evaluate(sequence):
        """评估序列质量"""
        # 颜色切换
        colors = [color_map[id] for id in sequence]
        switches = sum(1 for i in range(1, len(colors)) if colors[i] != colors[i-1])
        color_score = max(0, 100 - (switches / max(1, len(sequence)-1)) * 100)
        
        # 驱动比例平衡
        drives = [drive_map[id] for id in sequence]
        drive_deviation = 0
        for i in range(0, len(drives), 4):
            window = drives[i:i+4]
            two = window.count('两驱')
            four = window.count('四驱')
            if two + four > 0:
                deviation = abs(two - (two+four)/2) + abs(four - (two+four)/2)
                drive_deviation += deviation
        
        drive_score = max(0, 100 - min(100, drive_deviation))
        
        # 总得分
        total = color_score * 0.4 + drive_score * 0.3 + 85 * 0.2 + 95 * 0.1
        return total, switches, drive_deviation
    
    # 创建启发式序列
    electric = [id for id in cars if color_map[id] == '电动']
    combustion = [id for id in cars if color_map[id] == '燃动']
    
    two = [id for id in cars if drive_map[id] == '两驱']
    four = [id for id in cars if drive_map[id] == '四驱']
    
    # 方法1: 颜色优先
    color_grouped = electric + combustion
    
    # 方法2: 颜色优先+驱动平衡
    electric_two = [id for id in electric if drive_map[id] == '两驱']
    electric_four = [id for id in electric if drive_map[id] == '四驱']
    combustion_two = [id for id in combustion if drive_map[id] == '两驱']
    combustion_four = [id for id in combustion if drive_map[id] == '四驱']
    
    # 在每个颜色组内平衡驱动
    electric_balanced = []
    i = j = 0
    while i < len(electric_two) or j < len(electric_four):
        if i < len(electric_two):
            electric_balanced.append(electric_two[i])
            i += 1
        if j < len(electric_four):
            electric_balanced.append(electric_four[j])
            j += 1
    
    combustion_balanced = []
    i = j = 0
    while i < len(combustion_two) or j < len(combustion_four):
        if i < len(combustion_two):
            combustion_balanced.append(combustion_two[i])
            i += 1
        if j < len(combustion_four):
            combustion_balanced.append(combustion_four[j])
            j += 1
    
    optimized = electric_balanced + combustion_balanced
    
    # 比较所有方法
    methods = {
        '原始顺序': cars,
        '颜色分组': color_grouped,
        '颜色+驱动平衡': optimized
    }
    
    print("\n=== 优化方法比较 ===")
    results = []
    for name, seq in methods.items():
        score, switches, dev = evaluate(seq)
        results.append((name, score, switches, dev))
        print(f"{name}: 得分={score:.2f}, 颜色切换={switches}, 驱动偏差={dev:.1f}")
    
    # 选择最佳方案
    best_name, best_score, best_switches, best_dev = max(results, key=lambda x: x[1])
    best_sequence = methods[best_name]
    
    print(f"\n=== 最佳方案: {best_name} ===")
    print(f"总得分: {best_score:.2f}")
    print(f"颜色切换次数: {best_switches}")
    print(f"驱动比例偏差: {best_dev:.1f}")
    
    # 创建调度矩阵
    n_cols = max(32, (len(best_sequence) + 9) // 10)
    schedule = [[0] * n_cols for _ in range(10)]
    
    car_idx = 0
    for col in range(n_cols):
        for row in range(10):
            if car_idx < len(best_sequence):
                schedule[row][col] = best_sequence[car_idx]
                car_idx += 1
    
    # 保存结果
    pd.DataFrame(schedule).to_excel('result11.xlsx', index=False, header=False)
    
    sequence_data = []
    for car_id in best_sequence:
        row = df[df['sequence'] == car_id].iloc[0]
        sequence_data.append({
            '车辆顺序': int(car_id),
            '车型': str(row['config']),
            '动力': str(row['color']),
            '驱动': str(row['drive'])
        })
    
    pd.DataFrame(sequence_data).to_excel('result12.xlsx', index=False)
    
    # 最终统计
    colors = [color_map[id] for id in best_sequence]
    drives = [drive_map[id] for id in best_sequence]
    final_switches = sum(1 for i in range(1, len(colors)) if colors[i] != colors[i-1])
    
    print(f"\n=== 最终统计 ===")
    print(f"最终颜色切换次数: {final_switches}")
    print(f"电动车辆: {colors.count('电动')}辆")
    print(f"燃动车辆: {colors.count('燃动')}辆")
    print(f"两驱车辆: {drives.count('两驱')}辆")
    print(f"四驱车辆: {drives.count('四驱')}辆")
    print(f"结果已保存到 result11.xlsx 和 result12.xlsx")

if __name__ == "__main__":
    main()